#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
from PyQt6.QtWidgets import (QListWidget, QListWidgetItem, QFrame, QLabel, 
                           QVBoxLayout, QTableWidget, QTableWidgetItem, 
                           QHeaderView, QWidget, QAbstractItemView, QScrollArea,
                           QApplication, QSizePolicy)
from PyQt6.QtGui import QDrag, QFont, QColor, QPixmap
from PyQt6.QtCore import Qt, QMimeData, QByteArray, QPoint, pyqtSignal

class DraggableListWidget(QListWidget):
    """A list widget that supports drag and drop operations."""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setDragEnabled(True)
        self.setAcceptDrops(True)
        self.setDropIndicatorShown(True)
        self.setDefaultDropAction(Qt.DropAction.MoveAction)
        self.setSelectionMode(QAbstractItemView.SelectionMode.SingleSelection)
    
    def startDrag(self, supported_actions):
        """Start a drag operation for the selected items."""
        item = self.currentItem()
        if item:
            drag = QDrag(self)
            mime_data = QMimeData()
            
            # Get the class info from the item
            class_info = item.data(Qt.ItemDataRole.UserRole)
            
            # Set both text and custom MIME data
            mime_data.setText(self.format_class_info(class_info))
            mime_data.setData(
                "application/x-class-info", 
                QByteArray(str(class_info).encode('utf-8'))
            )
            
            drag.setMimeData(mime_data)
            drag.exec(Qt.DropAction.MoveAction, Qt.DropAction.MoveAction)
    
    def format_class_info(self, class_info):
        """Format class information for display."""
        if isinstance(class_info, str):
            return class_info
        return f"{class_info.get('college', '')} {class_info.get('major', '')} {class_info.get('grade', '')}级\n" \
               f"男生: {class_info.get('male', 0)} 女生: {class_info.get('female', 0)}"
    
    def dragEnterEvent(self, event):
        """Accept the event if it has a class info format."""
        if event.mimeData().hasFormat("application/x-class-info"):
            event.acceptProposedAction()
            self.setStyleSheet("""
                QTableWidget {
                    border: 2px dashed #2196F3;
                    background-color: #f5f9ff;
                }
                QTableWidget::item {
                    border: 1px solid #90caf9;
                    background-color: #e3f2fd;
                }
            """)
        else:
            event.ignore()
            
    def dragLeaveEvent(self, event):
        """Reset style when drag leaves."""
        self.setStyleSheet("")
        super().dragLeaveEvent(event)
    
    def dragMoveEvent(self, event):
        """Handle drag move event."""
        if event.mimeData().hasText() or event.mimeData().hasFormat("application/x-class-info"):
            event.acceptProposedAction()
        else:
            event.ignore()
    
    def dropEvent(self, event):
        """Handle drop event."""
        if event.source() == self:
            event.ignore()
            return
            
        if event.mimeData().hasText():
            event.acceptProposedAction()
            
            # Get the class info from the source
            class_info_text = event.mimeData().text()
            
            # Notify parent to handle the drop
            if hasattr(self.parent(), 'handle_class_return'):
                self.parent().handle_class_return(class_info_text)


class ClassCell(QFrame):
    """A widget that displays class information and supports drag and drop."""
    
    def __init__(self, class_info, parent=None):
        super().__init__(parent)
        self.class_info = class_info
        self.setAcceptDrops(True)
        self.setMinimumSize(150, 40)  # Set minimum size for better visibility
        self.init_ui()
    
    def init_ui(self):
        """Initialize the UI components."""
        self.setFrameShape(QFrame.Shape.StyledPanel)
        self.setLineWidth(1)
        self.update_theme(False)  # Initialize with light theme by default
        
        # Create layout and labels
        self.layout = QVBoxLayout(self)
        self.layout.setContentsMargins(4, 3, 4, 3)
        self.layout.setSpacing(2)
        
        # Get class info
        college = self.class_info.get('college', '')
        major = self.class_info.get('major', '')
        grade = self.class_info.get('grade', '')
        male = self.class_info.get('male', 0)
        female = self.class_info.get('female', 0)
        
        # Create labels for better control
        self.info_label = QLabel(f"{college} {major}")
        self.detail_label = QLabel(f"{grade}级  男:{male} 女:{female}")
        
        # Enable mouse tracking for better drag experience
        self.setMouseTracking(True)
        
        # Set font
        font = self.font()
        font.setPointSize(8)
        self.info_label.setFont(font)
        font.setBold(True)
        font.setPointSize(8)
        self.detail_label.setFont(font)
        
        # Set alignment and word wrap
        self.info_label.setWordWrap(True)
        self.detail_label.setWordWrap(True)
        
        # Add to layout
        self.layout.addWidget(self.info_label)
        self.layout.addWidget(self.detail_label)
        
        # Set fixed width to prevent taking full width
        self.setFixedWidth(180)  # Adjust width as needed

    def update_theme(self, is_dark):
        """Update the widget's appearance based on the current theme."""
        if is_dark:
            bg_color = "#2a2a2a"
            border_color = "#3a3a3a"
            text_color = "#ffffff"
            hover_color = "#3a3a3a"
            highlight_color = "#4a4a4a"
        else:
            bg_color = "#f0f5ff"  # Slightly blueish background
            border_color = "#cce0ff"
            text_color = "#333333"
            hover_color = "#e0ebff"
            highlight_color = "#d6e6ff"
        
        self.setStyleSheet(f"""
            ClassCell {{
                background-color: {bg_color};
                border: 1px solid {border_color};
                border-radius: 4px;
                margin: 1px;
            }}
            ClassCell:hover {{
                background-color: {hover_color};
                border-color: #99b3ff;
            }}
            QLabel {{
                color: {text_color};
                background: transparent;
                padding: 0px 2px;
                margin: 0px;
                line-height: 1.2;
            }}
        """)
    
    def mousePressEvent(self, event):
        """Handle mouse press event to start drag operation."""
        if event.button() == Qt.MouseButton.LeftButton:
            self.drag_start_position = event.pos()
            self.dragging = False
            # Visual feedback for the item being dragged
            self.setStyleSheet("""
                QFrame {
                    border: 2px solid #2196F3;
                    border-radius: 4px;
                    background-color: #e3f2fd;
                    padding: 2px;
                }
            """)
        super().mousePressEvent(event)
        
    def mouseReleaseEvent(self, event):
        """Reset style when mouse is released."""
        self.setStyleSheet("")
        super().mouseReleaseEvent(event)
    
    def mouseMoveEvent(self, event):
        """Handle mouse move event to start drag operation."""
        if not (event.buttons() & Qt.MouseButton.LeftButton) or hasattr(self, 'drag_start_position') is False:
            return
            
        if not self.dragging and \
           (event.pos() - self.drag_start_position).manhattanLength() < QApplication.startDragDistance():
            return
            
        self.dragging = True
        
        # Create drag object
        drag = QDrag(self)
        mime_data = QMimeData()
        
        # Create a simple preview
        preview = QLabel(self.format_class_info(), self)
        preview.setStyleSheet("""
            QLabel {
                background-color: #e1f5fe;
                border: 2px solid #2196F3;
                border-radius: 4px;
                padding: 4px 8px;
                color: #0d47a1;
                font-weight: bold;
            }
        """)
        preview.adjustSize()
        
        # Create pixmap for drag preview
        pixmap = QPixmap(preview.size())
        preview.render(pixmap)
        drag.setPixmap(pixmap)
        
        # Set mime data
        mime_data.setData("application/x-class-info", 
                         QByteArray(str(self.class_info).encode('utf-8')))
        drag.setMimeData(mime_data)
        
        # Start drag operation
        drag.setHotSpot(event.pos() - self.rect().topLeft())
        
        # Execute drag and handle result
        result = drag.exec(Qt.DropAction.MoveAction)
        if result == Qt.DropAction.MoveAction:
            self.hide()  # Hide the original widget if moved
            self.deleteLater()
    
    def format_class_info(self):
        """Format class information for display."""
        return f"{self.class_info.get('college', '')} {self.class_info.get('major', '')} {self.class_info.get('grade', '')}级\n" \
               f"男生: {self.class_info.get('male', 0)} 女生: {self.class_info.get('female', 0)}"


class ScheduleTable(QTableWidget):
    """A table widget that displays the class schedule and supports drag and drop."""
    
    # Signal emitted when the schedule is changed
    schedule_changed = pyqtSignal()
    
    def __init__(self, parent=None):
        super().__init__(3, 5, parent)  # 3 time slots, 5 days
        self.setAcceptDrops(True)
        self.setDragEnabled(True)
        self.setDragDropMode(QAbstractItemView.DragDropMode.DragDrop)
        self.setDefaultDropAction(Qt.DropAction.MoveAction)
        self.setDragDropOverwriteMode(False)
        
        # Set up table sizing
        self.setMinimumHeight(600)  # Set a minimum height for the table
        self.horizontalHeader().setDefaultSectionSize(220)  # Slightly wider for better text display
        self.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Fixed)
        self.verticalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Stretch)  # Make rows stretch to fill height
        self.verticalHeader().setMinimumSectionSize(100)  # Minimum height for each row
        self.setShowGrid(False)  # Hide grid lines for cleaner look
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)  # Allow table to expand
        
        # Enable drop indicator
        self.setDropIndicatorShown(True)
        
        # Set up headers
        self.setHorizontalHeaderLabels(["星期一", "星期二", "星期三", "星期四", "星期五"])
        self.setVerticalHeaderLabels(["第1-2节", "第3-4节", "第5-6节"])
        
        # Set up table properties
        self.setSelectionMode(QAbstractItemView.SelectionMode.SingleSelection)
        self.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        
        # Set up style
        self.setStyleSheet("""
            QTableWidget {
                gridline-color: #dee2e6;
                background-color: white;
                border: 1px solid #dee2e6;
                alternate-background-color: #f8f9fa;
            }
            QHeaderView::section {
                background-color: #f8f9fa;
                padding: 12px 8px;
                border: 1px solid #dee2e6;
                font-weight: bold;
                font-size: 12px;
            }
            QTableWidget::item {
                padding: 4px;
                border: 1px solid #dee2e6;
            }
            QTableCornerButton::section {
                background-color: #f8f9fa;
                border: 1px solid #dee2e6;
            }
        """)
        
        # Initialize cells
        for row in range(self.rowCount()):
            for col in range(self.columnCount()):
                item = QTableWidgetItem()
                item.classes = []  # Store class info
                item.widgets = []   # Store cell widgets
                self.setItem(row, col, item)
    
    def dragEnterEvent(self, event):
        """Handle drag enter event."""
        if event.mimeData().hasText() or event.mimeData().hasFormat("application/x-class-info"):
            event.acceptProposedAction()
        else:
            event.ignore()
    
    def dragMoveEvent(self, event):
        """Handle drag move events with visual feedback."""
        if event.mimeData().hasFormat("application/x-class-info"):
            # Highlight the cell under cursor
            pos = event.position().toPoint()
            item = self.itemAt(pos)
            if item:
                self.setCurrentItem(item)
                self.setStyleSheet("QTableWidget::item:selected { background-color: #e3f2fd; }")
            event.acceptProposedAction()
        else:
            event.ignore()
            return
            
        if event.mimeData().hasText():
            event.acceptProposedAction()
    
    def dropEvent(self, event):
        """Handle drop events for adding classes to time slots."""
        self.setStyleSheet("")  # Reset style
        
        # Only accept drops with our custom data
        if not event.mimeData().hasFormat("application/x-class-info"):
            event.ignore()
            return
        
        # Get the position where the drop occurred
        pos = event.position().toPoint()
        item = self.itemAt(pos)
        
        if not item:
            event.ignore()
            return
            
        try:
            # Get the class info from mime data
            data = event.mimeData().data("application/x-class-info")
            class_info = eval(data.data().decode('utf-8'))
            
            # Initialize widgets list if it doesn't exist
            if not hasattr(item, 'widgets'):
                item.widgets = []
            
            # Check if this class already exists in the target cell
            class_exists = any(hasattr(w, 'class_info') and w.class_info.get('id') == class_info.get('id') 
                             for w in item.widgets)
            
            if not class_exists:
                # Create and add new class cell
                class_cell = ClassCell(class_info, self)
                item.widgets.append(class_cell)
                
                # Update the cell display
                self.update_cell_widgets(item.row(), item.column())
                
                # Accept the drop
                event.setDropAction(Qt.DropAction.MoveAction)
                event.accept()
                
                # Force update the UI
                self.viewport().update()
                
                # Emit schedule changed signal to update statistics
                self.schedule_changed.emit()
            else:
                event.ignore()
                
        except Exception as e:
            print(f"Error in dropEvent: {e}")
            event.ignore()
    
    def add_class_to_cell(self, row, col, class_info):
        """Add a class to the specified cell."""
        if not (0 <= row < self.rowCount() and 0 <= col < self.columnCount()):
            return False
            
        item = self.item(row, col)
        if not item:
            return False
        
        # Check if class already exists in this cell
        for existing in item.classes:
            if (existing['college'] == class_info['college'] and 
                existing['major'] == class_info['major'] and 
                existing['grade'] == class_info['grade']):
                return False
        
        # Add class to the cell
        item.classes.append(class_info)
        
        # Create and add the class cell widget
        class_cell = ClassCell(class_info, self)
        item.widgets.append(class_cell)
        
        # Update the cell display
        self.update_cell_widgets(row, col)
        
        return True
    
    def remove_class_from_cell(self, row, col, class_info):
        """Remove a class from the specified cell."""
        if not (0 <= row < self.rowCount() and 0 <= col < self.columnCount()):
            return False
            
        item = self.item(row, col)
        if not item or not hasattr(item, 'classes'):
            return False
        
        # Find and remove the class
        for i, cls in enumerate(item.classes):
            if (cls['college'] == class_info['college'] and 
                cls['major'] == class_info['major'] and 
                cls['grade'] == class_info['grade']):
                
                item.classes.pop(i)
                
                # Remove the corresponding widget
                if i < len(item.widgets):
                    widget = item.widgets.pop(i)
                    widget.deleteLater()
                
                # Update the cell display
                self.update_cell_widgets(row, col)
                return True
        
        return False
    
    def update_cell_widgets(self, row, col):
        """Update the widgets in the specified cell."""
        try:
            item = self.item(row, col)
            if not item or not hasattr(item, 'widgets'):
                return
                
            # Clear existing widget
            widget = self.cellWidget(row, col)
            if widget:
                widget.deleteLater()
            
            if not item.widgets:  # No widgets to show
                return
            
            # Create a container widget for the cell
            container = QWidget()
            container.setAcceptDrops(True)
            
            # Create layout for the container
            layout = QVBoxLayout(container)
            layout.setContentsMargins(2, 2, 2, 2)
            layout.setSpacing(2)
            
            # Add all widgets to the layout
            for w in item.widgets:
                if w:  # Ensure widget exists
                    # Enable drag for all class cells
                    w.setAcceptDrops(True)
                    layout.addWidget(w)
            
            # Add stretch to push widgets to the top
            layout.addStretch()
            
            # Set the container as the cell widget
            self.setCellWidget(row, col, container)
            
            # Calculate required height based on number of classes
            visible_items = min(3, len(item.widgets))
            required_height = visible_items * 50  # ~50px per class cell
            
            # Set row height with min and max limits
            min_height = 50  # Minimum height for one class
            max_height = 180  # Maximum height to prevent taking too much space
            row_height = min(max(required_height, min_height), max_height)
            
            self.setRowHeight(row, row_height)
            
        except Exception as e:
            print(f"Error updating cell widgets at ({row}, {col}): {str(e)}")
